import socket
import cv2
import av
import struct
from io import BytesIO

BUFFER_SIZE = 1048576  # 1MB缓冲区
SERVER_IP = '0.0.0.0'
SERVER_PORT = 6001

def main():
    print(f"启动服务，等待连接...")
    # 初始化TCP服务
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    sock.bind((SERVER_IP, SERVER_PORT))
    sock.listen(1)
    
    # 创建H265解码上下文（启用硬件加速）
    codec = av.CodecContext.create('hevc', 'r')
    
    while True:
        conn, addr = sock.accept()
        print(f"连接来自: {addr}")
        
        # 接收配置包
        config_data = b''
        while len(config_data) < 12:
            chunk = conn.recv(12 - len(config_data))
            if not chunk:
                break
            config_data += chunk
        if len(config_data) != 12 or struct.unpack("<L", config_data[:4])[0] != 0xFFFFFFFF:
            print("配置包错误")
            conn.close()
            continue
        width, height = struct.unpack("<LL", config_data[4:12])
        print(f"新分辨率: {width}x{height}")
        codec = av.CodecContext.create('hevc', 'r')  # 重新初始化解码器
        
        buffer = BytesIO()
        try:
            while True:
                # 接收帧头
                header = b''
                while len(header) < 8:
                    chunk = conn.recv(8 - len(header))
                    if not chunk:
                        break
                    header += chunk
                if len(header) < 8:
                    break
                total_size, nalu_count = struct.unpack("<LL", header)
                
                # 接收帧数据
                received = 0
                buffer = BytesIO()
                while received < total_size:
                    chunk = conn.recv(min(BUFFER_SIZE, total_size - received))
                    if not chunk:
                        break
                    buffer.write(chunk)
                    received += len(chunk)

                
                # 分割NAL单元并解码
                buffer.seek(0)
                raw_data = buffer.getvalue()
                pos = 0
                
                # 根据MCU发送的nalu_count分割数据
                for _ in range(nalu_count):
                    # 假设每个NALU前有4字节长度头
                    nalu_size = struct.unpack("<L", raw_data[pos:pos+4])[0]
                    pos += 4
                    nalu_data = raw_data[pos:pos+nalu_size]
                    pos += nalu_size
                    
                    # 构造AVPacket进行解码
                    packet = av.packet.Packet(nalu_data)
                    try:
                        frames = codec.decode(packet)
                        for frame in frames:
                            if isinstance(frame, av.VideoFrame):
                                img = frame.to_ndarray(format='bgr24')
                                cv2.imshow('H265 Stream', img)
                                if cv2.waitKey(1) == 27:
                                    raise KeyboardInterrupt
                    except Exception as e:
                        print(f"解码错误: {e}")
                
                buffer = BytesIO()  # 重置缓冲区
                
        except (ConnectionResetError, KeyboardInterrupt):
            print("连接断开")
        finally:
            conn.close()
            cv2.destroyAllWindows()

if __name__ == "__main__":
    main()
